# -*- coding: UTF-8 -*-
# @File : chat_history.py
# @IDE     : VScode
# @Author  : zhonggc
# @Date    : 2024/05/27 11:46:38

import sqlite3
import uuid
from langchain_core.messages import HumanMessage, AIMessage, BaseMessage
import json
import datetime
from typing import Optional
import os


class ChatHistory:

    db_path = r".\src\ChatBot\DataBase\sqlite.db"

    def db_conn(self):
        conn = sqlite3.connect(self.db_path)
        return conn

    def save_chat_history(
        self, session_id: Optional[str] = None, messages: list[BaseMessage] = None
    ) -> None:
        """保存会话记录到数据库

        Arguments:
            session_id -- 会话唯一session_id
            messages -- 会话消息列表
        """
        try:
            if session_id is None:
                session_id = str(uuid.uuid1())
            conn = self.db_conn()
            message_list = [self.message2dict(message) for message in messages]
            if session_id in self.session_ids:
                # 更新会话以及更新时间
                sql_str = "UPDATE chat_history SET messages = ?, update_time = ? WHERE session_id = ?"
                conn.execute(
                    sql_str,
                    (
                        json.dumps(message_list, ensure_ascii=False),
                        str(datetime.datetime.now()),
                        session_id,
                    ),
                )
            else:
                # 新建会话
                sql_str = "INSERT INTO chat_history (session_id, title, messages, update_time) VALUES (?, ?, ?, ?)"
                title = "新建会话"
                conn.execute(
                    sql_str,
                    (
                        session_id,
                        title,
                        json.dumps(message_list, ensure_ascii=False),
                        str(datetime.datetime.now()),
                    ),
                )
        except Exception as e:
            print(f"保存会话记录失败: {e}")
        else:
            conn.commit()
        finally:
            conn.close()

    @property
    def session_ids(self):
        """查询会话列表: 按照更新时间倒序排列"""
        # 获取所有session_id
        try:
            conn = self.db_conn()
            # 按update_time倒序排列
            session_ids = conn.execute(
                "SELECT session_id FROM chat_history ORDER BY update_time DESC"
            ).fetchall()
            return [session_id[0] for session_id in session_ids]
        except Exception as e:
            print(f"查询会话列表失败: {e}")
            return []
        finally:
            conn.close()

    def pretty_session_dict(self):
        """查询会话列表: 按照更新时间倒序排列"""
        # 获取所有session_id: title
        try:
            conn = self.db_conn()
            # 按update_time倒序排列
            session_ids = conn.execute(
                "SELECT session_id, title, update_time FROM chat_history ORDER BY update_time DESC"
            ).fetchall()
            return [
                {
                    "session_id": session_id[0],
                    "title": session_id[1],
                    "update_time": session_id[2],
                }
                for session_id in session_ids
            ]
        except Exception as e:
            print(f"查询会话列表失败: {e}")
            return []
        finally:
            conn.close()

    def load_chat_history(self, session_id: str) -> list[BaseMessage]:
        """读取数据库中的会话记录

        Arguments:
            session_id -- 会话唯一session_id
        """
        try:
            conn = self.db_conn()
            sql_str = "SELECT messages FROM chat_history WHERE session_id = ?"
            messages = conn.execute(sql_str, (session_id,)).fetchone()
            if messages is not None:
                message_list = json.loads(messages[0])
                return [self.dict2message(message) for message in message_list]
            return []
        except Exception as e:
            print(f"读取会话记录失败: {e}")
            return []
        finally:
            conn.close()

    def delete_chat_history(self, session_id: str) -> None:
        """删除数据库中的会话记录

        Arguments:
            session_id -- 会话唯一session_id
        """
        try:
            conn = self.db_conn()
            sql_str = "DELETE FROM chat_history WHERE session_id = ?"
            conn.execute(sql_str, (session_id,))
        except Exception as e:
            print(f"删除会话记录失败: {e}")
        else:
            conn.commit()
        finally:
            conn.close()

    def message2dict(self, message: BaseMessage) -> dict:
        """将消息对象转换为字典

        Arguments:
            message -- 消息对象
        """
        return message.dict()

    def dict2message(self, message_dict: dict) -> HumanMessage | AIMessage:
        """将字典转换为消息对象

        Arguments:
            message_dict -- 消息字典
        """
        if message_dict.get("type") == "human":
            return HumanMessage(**message_dict)
        return AIMessage(**message_dict)
